{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data as data\n",
    "\n",
    "import random\n",
    "import numbers\n",
    "from PIL import Image, ImageMath\n",
    "import os\n",
    "import os.path\n",
    "import numpy as np\n",
    "import struct\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "import faiss\n",
    "import time\n",
    "\n",
    "import pcl\n",
    "from PIL import Image, ImageDraw\n",
    "import faiss\n",
    "\n",
    "import timeit\n",
    "\n",
    "%matplotlib qt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KNNBuilder_GPU:\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "        self.dimension = 3\n",
    "        \n",
    "        # we need only a StandardGpuResources per GPU\n",
    "        self.res = faiss.StandardGpuResources()\n",
    "        self.res.setTempMemoryFraction(0.1)\n",
    "        self.flat_config = faiss.GpuIndexFlatConfig()\n",
    "        self.flat_config.device = 0\n",
    "        \n",
    "    def build_nn_index(self, database):\n",
    "        '''\n",
    "        :param database: numpy array of Nx3\n",
    "        :return: Faiss index, in CPU\n",
    "        '''\n",
    "        index = faiss.GpuIndexFlatL2(self.res, self.dimension, self.flat_config)  # dimension is 3\n",
    "        index.add(database)\n",
    "        return index\n",
    "    \n",
    "    def search_nn(self, index, query, k):\n",
    "        '''\n",
    "        :param index: Faiss index\n",
    "        :param query: numpy array of Nx3\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        D, I = index.search(query, k)\n",
    "        return D, I\n",
    "    \n",
    "    def self_build_search(self, x):\n",
    "        '''\n",
    "\n",
    "        :param x: numpy array of Nxd\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        x = np.ascontiguousarray(x, dtype=np.float32)\n",
    "        index = self.build_nn_index(x)\n",
    "        D, I = self.search_nn(index, x, self.k)\n",
    "        return D, I\n",
    "    \n",
    "\n",
    "class KNNBuilder:\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "        self.dimension = 3\n",
    "\n",
    "    def build_nn_index(self, database):\n",
    "        '''\n",
    "        :param database: numpy array of Nx3\n",
    "        :return: Faiss index, in CPU\n",
    "        '''\n",
    "        index = faiss.IndexFlatL2(self.dimension)  # dimension is 3\n",
    "        index.add(database)\n",
    "        return index\n",
    "\n",
    "    def search_nn(self, index, query, k):\n",
    "        '''\n",
    "        :param index: Faiss index\n",
    "        :param query: numpy array of Nx3\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        D, I = index.search(query, k)\n",
    "        return D, I\n",
    "\n",
    "    def self_build_search(self, x):\n",
    "        '''\n",
    "\n",
    "        :param x: numpy array of Nxd\n",
    "        :return: D: numpy array of Nxk\n",
    "                 I: numpy array of Nxk\n",
    "        '''\n",
    "        x = np.ascontiguousarray(x, dtype=np.float32)\n",
    "        index = self.build_nn_index(x)\n",
    "        D, I = self.search_nn(index, x, self.k)\n",
    "        return D, I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_velodyne_bin(path):\n",
    "    '''\n",
    "    :param path:\n",
    "    :return: homography matrix of the point cloud, N*4\n",
    "    '''\n",
    "    pc_list = []\n",
    "    with open(path, 'rb') as f:\n",
    "        content = f.read()\n",
    "        pc_iter = struct.iter_unpack('ffff', content)\n",
    "        for idx, point in enumerate(pc_iter):\n",
    "            pc_list.append([point[0], point[1], point[2], point[3]])\n",
    "    return np.asarray(pc_list, dtype=np.float32)\n",
    "\n",
    "\n",
    "def read_calib(path):\n",
    "    '''\n",
    "    :param path: string pointed to the calib txt\n",
    "    :return: P0(3x4), P2(3x4), Tr(3x4)\n",
    "    '''\n",
    "    P0 = None\n",
    "    P2 = None\n",
    "    Tr = None\n",
    "    \n",
    "    f = open(path, 'r')\n",
    "    lines = [l.rstrip() for l in f.readlines()]\n",
    "    f.close()\n",
    "    for line in lines:\n",
    "        line_splitted = line.split()\n",
    "        if len(line_splitted) < 1:\n",
    "            pass\n",
    "        elif line_splitted[0] == 'P0:':\n",
    "            P0 = np.asarray(line_splitted[1:], dtype=np.float64).reshape(3,4)\n",
    "        elif line_splitted[0] == 'P2:':\n",
    "            P2 = np.asarray(line_splitted[1:], dtype=np.float64).reshape(3,4)\n",
    "        elif line_splitted[0] == 'R0_rect:':\n",
    "            R0_rect = np.asarray(line_splitted[1:], dtype=np.float64).reshape(3,3)\n",
    "        elif line_splitted[0] == 'Tr_velo_to_cam:' or line_splitted[0] == 'Tr:':\n",
    "            Tr = np.asarray(line_splitted[1:], dtype=np.float64).reshape(3,4)\n",
    "    \n",
    "    assert P0 is not None\n",
    "    assert P2 is not None\n",
    "    assert Tr is not None\n",
    "    \n",
    "    return (P0, P2, Tr)\n",
    "\n",
    "\n",
    "class PCSampler:\n",
    "    def __init__(self, leaf_size, minimum_pc_num):       \n",
    "        self.leaf_size = leaf_size\n",
    "        self.minimum_pc_num = minimum_pc_num\n",
    "    \n",
    "    def sample_pc(self, pc, leaf_size):\n",
    "        '''\n",
    "        :param pc: input numpy array of Nx3\n",
    "        :return: sampled_pc of Mx3\n",
    "        '''\n",
    "        cloud = pcl.PointCloud(pc)\n",
    "        sor = cloud.make_voxel_grid_filter()\n",
    "        sor.set_leaf_size(leaf_size, leaf_size, leaf_size)\n",
    "        cloud_filtered = sor.filter()\n",
    "        sampled_pc = np.asarray(cloud_filtered)\n",
    "        \n",
    "        return sampled_pc\n",
    "    \n",
    "    def sample_pc_wrapper(self, pc):\n",
    "        '''\n",
    "        ensure that the sampled pc is more than a certain amount\n",
    "        '''\n",
    "        retry_counter = 0\n",
    "        \n",
    "        sampled_pc = self.sample_pc(pc, self.leaf_size)\n",
    "        while sampled_pc.shape[0] < self.minimum_pc_num:\n",
    "            retry_counter += 1\n",
    "            leaf_size = self.leaf_size - 0.04*retry_counter\n",
    "            if leaf_size <= 0:\n",
    "                break\n",
    "            sampled_pc = self.sample_pc(pc, leaf_size)\n",
    "        \n",
    "        return sampled_pc\n",
    "    \n",
    "\n",
    "def transform_velo_to_camera(pc, Tr):\n",
    "    # pc: Nx4, x,y,z,reflectance\n",
    "    # Tr: 4x3\n",
    "    assert pc.shape[1] == 4\n",
    "    \n",
    "    pc_coord = pc[:, 0:3]  # Nx3\n",
    "    pc_coord_homo = np.concatenate((pc_coord, np.ones((pc_coord.shape[0], 1))), axis=1)  # Nx4\n",
    "    pc_coord_homo_transposed = np.transpose(pc_coord_homo)\n",
    "    pc_coord_cam = np.transpose(np.dot(Tr, pc_coord_homo_transposed))  # 3xN -> Nx3\n",
    "    \n",
    "    pc_cam = np.concatenate((pc_coord_cam, pc[:, 3:]), axis=1)\n",
    "    return pc_cam\n",
    "\n",
    "\n",
    "def project_cam_pc_to_img(pc, P):\n",
    "    # pc: Nx3, P: 3x4\n",
    "    assert pc.shape[1] == 3\n",
    "    \n",
    "    pc_homo = np.concatenate((pc, np.ones((pc.shape[0], 1))), axis=1)  # Nx4\n",
    "    pc_homo_transposed = np.transpose(pc_homo)  # 4xN\n",
    "    pc_projected = np.transpose(np.dot(P, pc_homo_transposed))  # Nx3\n",
    "    \n",
    "    pixels = pc_projected[:, 0:2] / pc_projected[:, 2:]\n",
    "    return pixels\n",
    "    \n",
    "    \n",
    "def crop_pc_within_img(pc, pixels):\n",
    "    # pc: Nx4, x, y, z, reflectance\n",
    "    # pixels: Nx2, x, y\n",
    "    assert pc.shape[1] == 4\n",
    "    \n",
    "    width, height = 1242, 375\n",
    "    x_valid_idx = np.logical_and(pixels[:, 0]>=0, pixels[:, 0]<=width)\n",
    "    y_valid_idx = np.logical_and(pixels[:, 1]>=0, pixels[:, 1]<=height)\n",
    "    pixel_valid_idx = np.logical_and(x_valid_idx, y_valid_idx)\n",
    "    range_valid_idx = pc[:, 2] > 0\n",
    "    valid_idx = np.logical_and(pixel_valid_idx, range_valid_idx)\n",
    "    \n",
    "    valid_pc = pc[valid_idx, :]\n",
    "    return valid_pc\n",
    "\n",
    "\n",
    "def axisEqual3D(ax):\n",
    "    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])\n",
    "    sz = extents[:,1] - extents[:,0]\n",
    "    centers = np.mean(extents, axis=1)\n",
    "    maxsize = max(abs(sz))\n",
    "    r = maxsize/2\n",
    "    for ctr, dim in zip(centers, 'xyz'):\n",
    "        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_from_bin(input_path, output_path, calib_path, leaf_size=0.2, minimum_pc_num=6000, \n",
    "                    radius_cutoff=90, outlier_filter=True):\n",
    "    pc_sampler = PCSampler(leaf_size, minimum_pc_num)\n",
    "    total_minimum_pc_num = 100000\n",
    "    \n",
    "    sample_num = len(os.listdir(input_path))\n",
    "    for i in range(sample_num):\n",
    "        \n",
    "        file = os.path.join(input_path, '%06d.bin'%i)\n",
    "        pc = read_velodyne_bin(file)\n",
    "        \n",
    "        # convert to camera coordinate, i.e., P0 coordinate\n",
    "        P0, P2, Tr = read_calib(calib_path)\n",
    "        pc_cam0 = transform_velo_to_camera(pc, Tr)  # Nx4\n",
    "        \n",
    "        # cutoff radius\n",
    "        valid_index = np.linalg.norm(pc_cam0[:, [0, 2]], axis=1) < radius_cutoff\n",
    "        pc_cam0 = np.ascontiguousarray(pc_cam0[valid_index, :], dtype=np.float32)\n",
    "        \n",
    "        # ------------------ verify the functions -------------------------------------------\n",
    "        # ------------------ verify the functions -------------------------------------------\n",
    "        \n",
    "        # down sample using pcl voxelgrid filter, \n",
    "        # \"all the points present will be approximated (i.e., downsampled) with their centroid\"\n",
    "        sampled_pc_cam0 = pc_sampler.sample_pc_wrapper(pc_cam0[:, 0:3].astype(np.float32))\n",
    "        \n",
    "        # outlier filter\n",
    "        if outlier_filter == True:\n",
    "#             start_t = timeit.default_timer()\n",
    "            fil = pcl.PointCloud(sampled_pc_cam0).make_statistical_outlier_filter()\n",
    "            fil.set_mean_k(50)\n",
    "            fil.set_std_dev_mul_thresh(2)\n",
    "            sampled_pc_cam0 = np.asarray(fil.filter())\n",
    "#             stop_t = timeit.default_timer()\n",
    "#             print('outlier filter: %f' % (stop_t - start_t))\n",
    "        \n",
    "        # KNN to get the reflectance, and measure the error introduced by voxel filter\n",
    "        knn = KNNBuilder_GPU(k=1)\n",
    "        index = knn.build_nn_index(np.ascontiguousarray(pc_cam0[:, 0:3], dtype=np.float32))\n",
    "        D, I = knn.search_nn(index, sampled_pc_cam0.astype(np.float32), k=1)\n",
    "        \n",
    "        reflectance = pc_cam0[:, 3:][I[:,0]]\n",
    "        sampled_pc_cam0 = np.concatenate((sampled_pc_cam0, reflectance), axis=1)\n",
    "        # print(sampled_pc_cam0)\n",
    "        # print(pc_cam0[I[2,0]])\n",
    "        \n",
    "        # so far it seems the sampled_pc_cam0 is ready for saving\n",
    "        if outlier_filter == False:\n",
    "            if not os.path.isdir(os.path.join(output_path, 'np_%.2f_%d_r%d'%(leaf_size, minimum_pc_num, radius_cutoff))):\n",
    "                os.makedirs(os.path.join(output_path, 'np_%.2f_%d_r%d'%(leaf_size, minimum_pc_num, radius_cutoff)))\n",
    "            np.save(os.path.join(output_path, 'np_%.2f_%d_r%d'%(leaf_size, minimum_pc_num, radius_cutoff), '%06d.npy'%i), sampled_pc_cam0)\n",
    "        else:\n",
    "            if not os.path.isdir(os.path.join(output_path, 'np_%.2f_%d_r%d_filtered'%(leaf_size, minimum_pc_num, radius_cutoff))):\n",
    "                os.makedirs(os.path.join(output_path, 'np_%.2f_%d_r%d_filtered'%(leaf_size, minimum_pc_num, radius_cutoff)))\n",
    "            np.save(os.path.join(output_path, 'np_%.2f_%d_r%d_filtered'%(leaf_size, minimum_pc_num, radius_cutoff), '%06d.npy'%i), sampled_pc_cam0)\n",
    "        \n",
    "        # debug\n",
    "#         if True:\n",
    "#             print(pc.shape)\n",
    "#             print(pc_cam0.shape)\n",
    "#             print(sampled_pc_cam0.shape)\n",
    "#             print('mean distance from sampled pt to original pt: %f'%np.mean(D))\n",
    "#             print('max distance from sampled pt to original pt: %f'%np.amax(D))\n",
    "#             print('distance std sigma from sampled pt to original pt: %f'%np.std(D))\n",
    "#             fig = plt.figure()\n",
    "#             ax = Axes3D(fig)\n",
    "#             ax.scatter(sampled_pc_cam0[:,0].tolist(), sampled_pc_cam0[:,1].tolist(), sampled_pc_cam0[:,2].tolist(), s=0.1, c=[0.5,0.5,0.5])\n",
    "#             axisEqual3D(ax)\n",
    "\n",
    "#             fig = plt.figure()\n",
    "#             ax = Axes3D(fig)\n",
    "#             ax.scatter(pc_cam0[:,0].tolist(), pc_cam0[:,1].tolist(), pc_cam0[:,2].tolist(), s=0.1, c=[0.5,0.5,0.5])\n",
    "#             axisEqual3D(ax)\n",
    "\n",
    "#             plt.ion()\n",
    "#             plt.show()\n",
    "            \n",
    "#             break\n",
    "                    \n",
    "        \n",
    "        if sampled_pc_cam0.shape[0] < total_minimum_pc_num:\n",
    "            total_minimum_pc_num = sampled_pc_cam0.shape[0]\n",
    "        \n",
    "        if i % 100 == 0:\n",
    "            print('iteration: %d, total_minimum_pc_num: %d' % (i, total_minimum_pc_num))\n",
    "        \n",
    "        \n",
    "    return total_minimum_pc_num\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "minimum_pc_num = 100000\n",
    "for seq in range(11):\n",
    "    begin_t = time.time()\n",
    "    pc_num = sample_from_bin('/ssd/dataset/odometry/data_odometry_velodyne/sequences/%02d/velodyne'%(seq),\n",
    "                             '/ssd/dataset/odometry/data_odometry_velodyne/numpy/%02d'%(seq),\n",
    "                             '/ssd/dataset/odometry/calib/%02d/calib.txt'%(seq),\n",
    "                             leaf_size=0.2, minimum_pc_num=20480, radius_cutoff=50, outlier_filter=False)\n",
    "    print(time.time()-begin_t)\n",
    "    \n",
    "    if pc_num < minimum_pc_num:\n",
    "        minimum_pc_num = pc_num\n",
    "\n",
    "    print('minimum_pc_num: %d' % minimum_pc_num)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minimum_pc_num = 100000\n",
    "\n",
    "seq = 10\n",
    "\n",
    "begin_t = time.time()\n",
    "pc_num = sample_from_bin('/ssd/dataset/odometry/data_odometry_velodyne/sequences/%02d/velodyne'%(seq),\n",
    "                         '/ssd/dataset/odometry/data_odometry_velodyne/numpy/%02d'%(seq),\n",
    "                         '/ssd/dataset/odometry/calib/%02d/calib.txt'%(seq),\n",
    "                         leaf_size=0.2, minimum_pc_num=20480, radius_cutoff=90, outlier_filter=False)\n",
    "print(time.time()-begin_t)\n",
    "\n",
    "if pc_num < minimum_pc_num:\n",
    "    minimum_pc_num = pc_num\n",
    "\n",
    "print('minimum_pc_num: %d' % minimum_pc_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# get surface normal\n",
    "def Surface_normals(cloud):\n",
    "    ne = cloud.make_NormalEstimation()\n",
    "    tree = cloud.make_kdtree()\n",
    "    ne.set_SearchMethod(tree)\n",
    "#     ne.set_RadiusSearch(2)\n",
    "    ne.set_KSearch(9)\n",
    "    cloud_normals = ne.compute()\n",
    "    return cloud_normals\n",
    "\n",
    "\n",
    "input_path = '/ssd/dataset/odometry/data_odometry_velodyne/numpy'\n",
    "input_folder = 'np_0.20_20480_r90'\n",
    "output_folder = 'np_0.20_20480_r90_sn'\n",
    "for seq in range(11):\n",
    "    sample_num = len(os.listdir(os.path.join(input_path, '%02d'%seq, input_folder)))\n",
    "    for i in range(sample_num):\n",
    "        pc_np = np.load(os.path.join(input_path, '%02d'%seq, input_folder, '%06d.npy'%i))  # Nx4, x,y,z,ref\n",
    "        cloud = pcl.PointCloud(pc_np[:, 0:3])\n",
    "        \n",
    "        sn = Surface_normals(cloud)\n",
    "        sn_np = (np.asarray(sn.to_array()))  # Nx4, nx,ny,nz,curvature\n",
    "        \n",
    "        output_np = np.concatenate((pc_np[:, 0:3], sn_np, pc_np[:, 3:]), axis=1)\n",
    "        \n",
    "        if not os.path.isdir(os.path.join(input_path, '%02d'%seq, output_folder)):\n",
    "            os.makedirs(os.path.join(input_path, '%02d'%seq, output_folder))\n",
    "        np.save(os.path.join(input_path, '%02d'%seq, output_folder, '%06d.npy'%i), output_np)\n",
    "        \n",
    "        if i%100==0:\n",
    "            print('seq %d - %d' % (seq, i))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_path = '/ssd/dataset/odometry/data_odometry_velodyne/numpy'\n",
    "input_folder = 'np_0.20_20480_r90_filtered'\n",
    "output_folder = 'np_0.20_20480_r90_filtered_sn'\n",
    "\n",
    "pc_np = np.load(os.path.join(input_path, '06', output_folder, '000100.npy'))\n",
    "print(pc_np.shape)\n",
    "\n",
    "# mean of three axis\n",
    "print('max of three axis')\n",
    "print(np.max(pc_np[:, 0:3], axis=0))\n",
    "print('min of three axis')\n",
    "print(np.min(pc_np[:, 0:3], axis=0))\n",
    "\n",
    "print(np.linalg.norm(pc_np[:, 3:6], axis=1))\n",
    "# reflectance\n",
    "print(np.mean(pc_np[:, 7]))\n",
    "print(np.max(pc_np[:, 7]))\n",
    "print(np.min(pc_np[:, 7]))\n",
    "\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = Axes3D(fig)\n",
    "ax.scatter(pc_np[:,0].tolist(), pc_np[:,1].tolist(), pc_np[:,2].tolist(), s=0.1, c=[0.5,0.5,0.5])\n",
    "axisEqual3D(ax)\n",
    "\n",
    "plt.ion()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
